#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Delta 机器人工作空间点云可视化（快速修法版本）
- 采用 ProcessPoolExecutor.map + chunksize，避免对同一迭代器并发 islice 的问题
- 其余逻辑基本保持不变
"""

import os
import argparse
import itertools
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider
from concurrent.futures import ProcessPoolExecutor


# ---------------- Delta 正解 -----------------
def delta_forward_kinematics(theta1, theta2, theta3,
                             R=90.0, r=41.57, L=85.0, l=140.0):
    """
    输入角度单位：rad
    输出坐标单位：mm
    """
    a1 = np.array([ R,            0,               0])
    a2 = np.array([-R/2,  R*np.sqrt(3)/2,          0])
    a3 = np.array([-R/2, -R*np.sqrt(3)/2,          0])

    a1b1 = np.array([ L*np.cos(theta1),                 0,                L*np.sin(theta1)])
    a2b2 = np.array([-L/2*np.cos(theta2),  L*np.cos(theta2)*np.sqrt(3)/2, L*np.sin(theta2)])
    a3b3 = np.array([-L/2*np.cos(theta3), -L*np.cos(theta3)*np.sqrt(3)/2, L*np.sin(theta3)])

    c1o2 = np.array([-r,                 0,                0])
    c2o2 = np.array([ r/2, -r*np.sqrt(3)/2,                0])
    c3o2 = np.array([ r/2,  r*np.sqrt(3)/2,                0])

    o1d1 = a1 + a1b1 + c1o2
    o1d2 = a2 + a2b2 + c2o2
    o1d3 = a3 + a3b3 + c3o2

    o1f  = (o1d1 + o1d2) / 2.0

    d1d2 = o1d2 - o1d1
    d1d3 = o1d3 - o1d1
    d2d3 = o1d3 - o1d2

    a = np.linalg.norm(d1d2)
    b = np.linalg.norm(d1d3)
    c = np.linalg.norm(d2d3)
    p = (a + b + c) / 2.0
    S = np.sqrt(p*(p-a)*(p-b)*(p-c))

    normd1e  = a*b*c/(4.0*S)
    normfe   = np.sqrt(normd1e**2 - (a**2)/4.0)
    dirfe1   = np.cross(np.cross(d1d2, d1d3), d1d2) / (a*a*b)   # 方向
    dirfe    = dirfe1 / np.linalg.norm(dirfe1)
    fe       = normfe * dirfe
    o1e      = o1f + fe

    normd1o2 = l
    normeo2  = np.sqrt(normd1o2**2 - normd1e**2)
    direo21  = np.cross(d1d2, d1d3) / (a*b)
    direo2   = direo21 / np.linalg.norm(direo21)
    eo2      = normeo2 * direo2
    o1o2     = o1e + eo2
    return o1o2  # x,y,z
# --------------------------------------------------


# ---------------- Delta 逆解 -----------------
def delta_inverse_kinematics(x, y, z, 
                            R=90.0, r=41.57, L=85.0, l=140.0):
    """
    输入坐标单位：mm
    输出角度单位：rad
    """

    m = x**2 + y**2 + z**2 + (R-r)**2 + L**2 - l**2
    
    A = np.array([
        (m - 2*x*(R-r))/(2*L) - (R - r - x),
        (m + (R-r)*(x - np.sqrt(3)*y))/L - 2*(R - r) - (x - np.sqrt(3)*y),
        (m + (R-r)*(x + np.sqrt(3)*y))/L - 2*(R - r) - (x + np.sqrt(3)*y)
    ])
    
    B = np.array([2*z, 4*z, 4*z])
    
    C = np.array([
        (m - 2*x*(R-r))/(2*L) + (R - r - x),
        (m + (R-r)*(x - np.sqrt(3)*y))/L + 2*(R - r) + (x - np.sqrt(3)*y),
        (m + (R-r)*(x + np.sqrt(3)*y))/L + 2*(R - r) + (x + np.sqrt(3)*y)
    ])
    
    # 求解二次方程
    theta1 = 2 * np.arctan((-B[0] - np.sqrt(B[0]**2 - 4*A[0]*C[0])) / (2*A[0]))
    theta2 = 2 * np.arctan((-B[1] - np.sqrt(B[1]**2 - 4*A[1]*C[1])) / (2*A[1]))
    theta3 = 2 * np.arctan((-B[2] - np.sqrt(B[2]**2 - 4*A[2]*C[2])) / (2*A[2]))
    
    return theta1, theta2, theta3
# --------------------------------------------------

# ---------- 参数解析 ----------
parser = argparse.ArgumentParser(description='Delta机器人工作空间可视化')
parser.add_argument('--R', type=float, default=72, help='静平台半径（默认：72）')
parser.add_argument('--r', type=float, default=36.15, help='动平台半径（默认：36.15）')
parser.add_argument('--L', type=float, default=200, help='驱动杆长度（默认：200）')
parser.add_argument('--l', type=float, default=230, help='平行四边形边长（默认：230）')
args = parser.parse_args()

# 使用命令行参数或默认值
R, r, L, l = args.R, args.r, args.L, args.l
angles_deg = np.arange(1, 136, 2)                # 角度采样
chunksize = 4096                                 # map 的块大小
tolerance = 5                                    # Z 切片厚度（mm）


# ---------- 单点计算 ----------
def calculate_point(angles_triplet):
    i, j, k = angles_triplet
    th1 = np.deg2rad(k - 46)
    th2 = np.deg2rad(j - 46)
    th3 = np.deg2rad(i - 46)
    x, y, z = delta_forward_kinematics(th1, th2, th3, R, r, L, l)
    return [x, y, -z]


def compute_points():
    combos = itertools.product(angles_deg, repeat=3)  # 共有 N^3 个组合
    with ProcessPoolExecutor(max_workers=os.cpu_count()) as ex:
        # ex.map 返回的是惰性迭代器，这里一次性收集为 list 再转成 ndarray
        pts_iter = ex.map(calculate_point, combos, chunksize=chunksize)
        pts = np.array(list(pts_iter), dtype=float)
    return pts


# ---------- 主程序 ---------- 
def main():
    pts = compute_points()
    
    # 数据清洗
    mask = ~np.isnan(pts).any(axis=1) & np.isfinite(pts).all(axis=1)
    pts = pts[mask]
    assert len(pts) > 0, "No valid points found after filtering NaN/Inf values"

    # ----------- 作图 -----------
    fig = plt.figure(figsize=(12, 6))
    
    # 3D 图
    ax3d = fig.add_subplot(121, projection='3d')
    ax3d.scatter(pts[:, 0], pts[:, 1], pts[:, 2], s=1)
    ax3d.set_box_aspect([1, 1, 1])
    ax3d.set_xlabel('X (mm)')
    ax3d.set_ylabel('Y (mm)')
    ax3d.set_zlabel('Z (mm)')
    ax3d.set_title('Workspace Point Cloud')

    # 2D 切片图
    ax2d = fig.add_subplot(122)
    ax2d.set_aspect('equal')
    ax2d.set_title('XY Plane Slice')
    ax2d.set_xlabel('X (mm)')
    ax2d.set_ylabel('Y (mm)')
    
    # 添加角度显示文本框
    angle_text = ax2d.text(0.05, 0.05, '', 
                          transform=ax2d.transAxes, 
                          bbox=dict(facecolor='white', alpha=0.8))

    # 滑块
    slider_ax = fig.add_axes([0.25, 0.08, 0.5, 0.03])
    z_min, z_max = pts[:, 2].min(), pts[:, 2].max()
    z_slider = Slider(slider_ax, 'Z Value', z_min, z_max, valinit=(z_min + z_max) / 2)

    # 初始切片
    current_z = (z_min + z_max) / 2
    slice_mask = (pts[:, 2] > current_z - tolerance) & (pts[:, 2] < current_z + tolerance)
    scatter2d = ax2d.scatter(pts[slice_mask, 0], pts[slice_mask, 1], s=1)

    def update(val):
        cz = z_slider.val
        sm = (pts[:, 2] > cz - tolerance) & (pts[:, 2] < cz + tolerance)
        scatter2d.set_offsets(pts[sm, :2])
        fig.canvas.draw_idle()

    # 鼠标点击事件处理
    def on_mouse_click(event):  # 重命名函数
        if event.inaxes == ax2d and event.button == 1:  # 添加左键点击判断
            try:
                x, y = event.xdata, event.ydata
                z = z_slider.val
                
                # 计算逆解
                theta1, theta2, theta3 = delta_inverse_kinematics(x, y, z, R, r, L, l)
                
                # 更新显示文本
                text = f'点击位置：\nX: {x:.1f} mm\nY: {y:.1f} mm\nZ: {z:.1f} mm\n\n' \
                       f'θ1: {np.rad2deg(theta1):.1f}°\nθ2: {np.rad2deg(theta2):.1f}°\nθ3: {np.rad2deg(theta3):.1f}°'
                angle_text.set_text(text)
            except Exception as e:
                angle_text.set_text(f'无效位置\n{str(e)}')
            fig.canvas.draw_idle()

    # 绑定点击事件（替换原来的移动事件）
    fig.canvas.mpl_connect('button_press_event', on_mouse_click)  # 修改事件类型
    z_slider.on_changed(update)

    plt.tight_layout(rect=[0, 0.12, 1, 1])
    plt.show()


if __name__ == '__main__':
    main()